import numpy as np
from PyCmpltrtok.common import *
import cv2 as cv
from PyCmpltrtok.common_opencv import *


def postprocess(output_arr, img_=None, seed=None):
    sep('postprocess')

    # category names
    label_names_path = '/var/asuspei/my_svn/c_darknet/darknet-debug-cpu/data/voc.names'
    with open(label_names_path, 'r') as f:
        names = f.readlines()
    names = [name[:-2] if name[-2] == '\r' else name[:-1] for name in names]
    names = np.array(names)
    print('names', names)

    check_np_detailed(output_arr, 'output_arr')
    output_arr = np.squeeze(output_arr, (0, 1))
    check_np_detailed(output_arr, 'output_arr')
    uniq = np.unique(output_arr)
    print('unique', uniq)
    objs = uniq[uniq > 0] -1
    print(names[objs])

    n_cls = len(names) + 1
    if 0:
        lut = []
        if seed is None:
            seed = 666
        np.random.seed(seed)
        print('seed', seed)
        for i in range(n_cls):
            lut.append(rand_color())
        lut = np.array(lut, dtype=np.uint8)
        lut[0] = (0, 0, 0)
    else:
        lut = rand_palette(n_cls)
    lut = lut[:, ::-1]
    print(lut)

    visual = lut[output_arr]
    check_np(visual, 'visual')
    if img_ is not None:
        visual = np.concatenate([img_, visual], axis=1)
    resized = imzoom2fit_rect(visual, (1600, 800))
    cv.imshow('visual', resized)
    H, W = 400, 300
    canvas = np.zeros((H, W, 3), dtype=np.uint8)
    base = 0
    for idx in uniq:
        color = lut[idx].tolist()
        cv.rectangle(canvas, (0, base), (20, base + 20), color, cv.FILLED)
        cv.putText(canvas, names[idx - 1], (30, base + 20), cv.FONT_HERSHEY_PLAIN, 1, color, 1)
        base += 30
    cv.imshow('legend', canvas)
    cv.waitKey()
    cv.destroyAllWindows()